#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
洪水预报智能代理使用示例

展示如何使用新的LangChain代理系统进行参数收集和验证
"""

import asyncio
import json
from typing import Dict, Any
from loguru import logger

# 导入代理模块
from __init__ import create_agent, FloodForecastAgent

# 配置日志
logger.add("logs/agent_example.log", rotation="1 day", retention="7 days")

def example_1_basic_usage():
    """示例1：基本使用方法"""
    print("\n=== 示例1：基本使用方法 ===")
    
    # 配置LLM
    llm_config = {
        "type": "ollama",
        "model": "qwq:latest",
        "temperature": 0.1,
        "base_url": "http://10.48.0.81:11434"
    }
    
    # 创建代理
    agent = create_agent(llm_config)
    
    # 模拟用户对话
    user_inputs = [
        "我需要对大伙房水库进行洪水预报",
        "使用新安江模型",
        "从2024-01-15 08:00开始预报",
        "预报48小时，降雨预见期12小时",
        "目标变量是流量"
    ]
    
    for i, user_input in enumerate(user_inputs, 1):
        print(f"\n第{i}轮对话:")
        print(f"用户: {user_input}")
        
        response = agent.process_user_input(user_input)
        
        print(f"助手: {response['response']}")
        print(f"当前阶段: {response['current_stage']}")
        print(f"已收集参数: {json.dumps(response['collected_parameters'], ensure_ascii=False, indent=2)}")
        print(f"缺失参数: {response['missing_parameters']}")
        print(f"完成状态: {response['is_complete']}")
        print(f"置信度: {response['confidence_score']:.2f}")
        
        if response['is_complete']:
            print("\n✅ 参数收集完成！")
            break
    
    # 获取最终摘要
    summary = agent.get_session_summary()
    print(f"\n会话摘要: {json.dumps(summary, ensure_ascii=False, indent=2)}")

def example_2_error_handling():
    """示例2：错误处理和参数修正"""
    print("\n=== 示例2：错误处理和参数修正 ===")
    
    llm_config = {
        "type": "ollama",
        "model": "qwq:latest",
        "temperature": 0.1,
        "base_url": "http://10.48.0.81:11434"
    }
    
    agent = create_agent(llm_config)
    
    # 包含错误的用户输入
    error_inputs = [
        "我要预报某个水库的洪水",  # 缺少具体水库名称
        "水库名称是不存在的水库",  # 无效水库名称
        "预报时间是昨天",  # 无效时间
        "预报1000小时",  # 超出合理范围
        "修正为大伙房水库，新安江模型，2024-01-15 08:00开始，预报48小时，降雨预见期12小时，目标变量流量"
    ]
    
    for i, user_input in enumerate(error_inputs, 1):
        print(f"\n第{i}轮对话:")
        print(f"用户: {user_input}")
        
        response = agent.process_user_input(user_input)
        
        print(f"助手: {response['response']}")
        
        if response.get('validation_result'):
            validation = response['validation_result']
            if validation.get('validation_errors'):
                print(f"验证错误: {json.dumps(validation['validation_errors'], ensure_ascii=False, indent=2)}")
            if validation.get('validation_warnings'):
                print(f"验证警告: {json.dumps(validation['validation_warnings'], ensure_ascii=False, indent=2)}")
        
        if response['is_complete']:
            print("\n✅ 参数收集完成！")
            break

def example_3_session_management():
    """示例3：会话管理"""
    print("\n=== 示例3：会话管理 ===")
    
    llm_config = {
        "type": "ollama",
        "model": "qwq:latest",
        "temperature": 0.1,
        "base_url": "http://10.48.0.81:11434"
    }
    
    agent = create_agent(llm_config)
    
    # 第一次对话
    print("\n--- 第一次对话 ---")
    response1 = agent.process_user_input("我需要对三峡水库进行洪水预报")
    print(f"助手: {response1['response']}")
    
    # 获取会话状态
    summary1 = agent.get_session_summary()
    print(f"会话状态: 阶段={summary1['current_stage']}, 参数数量={len(summary1['collected_parameters'])}")
    
    # 继续对话
    print("\n--- 继续对话 ---")
    response2 = agent.process_user_input("使用LSTM模型，从明天上午9点开始")
    print(f"助手: {response2['response']}")
    
    # 重置会话
    print("\n--- 重置会话 ---")
    agent.reset_session()
    summary_reset = agent.get_session_summary()
    print(f"重置后状态: 阶段={summary_reset['current_stage']}, 参数数量={len(summary_reset['collected_parameters'])}")
    
    # 重新开始
    print("\n--- 重新开始 ---")
    response3 = agent.process_user_input("我要预报小浪底水库的流量")
    print(f"助手: {response3['response']}")

def example_4_custom_configuration():
    """示例4：自定义配置"""
    print("\n=== 示例4：自定义配置 ===")
    
    # 自定义LLM配置
    custom_llm_config = {
        "type": "ollama",
        "model": "qwq:latest",
        "temperature": 0.05,  # 更低的温度，更确定性的回答
        "max_tokens": 1500,
        "base_url": "http://10.48.0.81:11434"
    }
    
    # 自定义代理配置
    custom_agent_config = {
        "max_iterations": None,  # 无迭代次数限制
        "max_execution_time": None  # 无执行时间限制
    }
    
    # 创建自定义代理
    agent = FloodForecastAgent(
        llm_config=custom_llm_config,
        **custom_agent_config
    )
    
    print("使用自定义配置创建的代理:")
    print(f"- 模型: {custom_llm_config['model']}")
    print(f"- 温度: {custom_llm_config['temperature']}")
    print(f"- 最大迭代: {custom_agent_config['max_iterations'] or '无限制'}")
    print(f"- 最大执行时间: {custom_agent_config['max_execution_time'] or '无限制'}")
    
    # 测试对话
    response = agent.process_user_input("我需要一个复杂的洪水预报方案")
    print(f"\n助手: {response['response']}")

def example_5_batch_processing():
    """示例5：批量处理"""
    print("\n=== 示例5：批量处理 ===")
    
    llm_config = {
        "type": "ollama",
        "model": "qwq:latest",
        "temperature": 0.1,
        "base_url": "http://10.48.0.81:11434"
    }
    
    # 批量处理多个预报请求
    forecast_requests = [
        {
            "name": "请求1",
            "inputs": [
                "大伙房水库洪水预报",
                "新安江模型，2024-01-15 08:00，48小时，12小时降雨预见期，流量"
            ]
        },
        {
            "name": "请求2", 
            "inputs": [
                "三峡水库水位预报",
                "LSTM模型，明天上午9点开始，72小时预报，24小时降雨预见期"
            ]
        }
    ]
    
    results = []
    
    for request in forecast_requests:
        print(f"\n处理 {request['name']}:")
        
        # 为每个请求创建新的代理实例
        agent = create_agent(llm_config)
        
        for user_input in request['inputs']:
            print(f"  输入: {user_input}")
            response = agent.process_user_input(user_input)
            print(f"  响应: {response['response'][:100]}...")
        
        # 收集结果
        final_summary = agent.get_session_summary()
        results.append({
            "request_name": request['name'],
            "is_complete": final_summary['is_complete'],
            "collected_parameters": final_summary['collected_parameters'],
            "confidence_score": final_summary['confidence_score']
        })
    
    print("\n批量处理结果:")
    for result in results:
        print(f"- {result['request_name']}: 完成={result['is_complete']}, 置信度={result['confidence_score']:.2f}")

async def example_6_async_usage():
    """示例6：异步使用（模拟）"""
    print("\n=== 示例6：异步使用模拟 ===")
    
    llm_config = {
        "type": "ollama",
        "model": "qwq:latest",
        "temperature": 0.1,
        "base_url": "http://10.48.0.81:11434"
    }
    
    # 模拟异步处理多个会话
    async def process_session(session_name: str, inputs: list):
        print(f"\n开始处理会话: {session_name}")
        agent = create_agent(llm_config)
        
        for user_input in inputs:
            # 模拟异步延迟
            await asyncio.sleep(0.1)
            response = agent.process_user_input(user_input)
            print(f"  {session_name}: {response['response'][:50]}...")
        
        summary = agent.get_session_summary()
        print(f"  {session_name} 完成: {summary['is_complete']}")
        return summary
    
    # 并发处理多个会话
    sessions = [
        ("会话A", ["大伙房水库预报", "新安江模型，明天8点，48小时，流量"]),
        ("会话B", ["三峡水库预报", "LSTM模型，后天9点，72小时，水位"]),
        ("会话C", ["小浪底水库预报", "HBV模型，今天下午，24小时，入库流量"])
    ]
    
    tasks = [process_session(name, inputs) for name, inputs in sessions]
    results = await asyncio.gather(*tasks)
    
    print("\n所有会话处理完成:")
    for i, result in enumerate(results):
        print(f"  会话{chr(65+i)}: 完成={result['is_complete']}, 置信度={result['confidence_score']:.2f}")

def main():
    """主函数 - 运行所有示例"""
    print("洪水预报智能代理使用示例")
    print("=" * 50)
    
    try:
        # 运行同步示例
        example_1_basic_usage()
        example_2_error_handling()
        example_3_session_management()
        example_4_custom_configuration()
        example_5_batch_processing()
        
        # 运行异步示例
        print("\n运行异步示例...")
        asyncio.run(example_6_async_usage())
        
        print("\n\n🎉 所有示例运行完成！")
        
    except Exception as e:
        logger.error(f"示例运行失败: {str(e)}")
        print(f"❌ 示例运行失败: {str(e)}")
        print("请检查LLM配置和网络连接")

if __name__ == "__main__":
    main()